#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2019 crane <crane@gosun>
#
# Distributed under terms of the MIT license.

"""
https://www.lintcode.com/problem/best-time-to-buy-and-sell-stock-ii/description

https://www.imooc.com/article/28425
"""

class Solution:
    """
    @param prices: Given an integer array
    @return: Maximum profit
    """
    def maxProfit(self, prices):
        # return self.method1(prices)
        return self.method2(prices)

    def method2(self, prices):
        """
        HINT: 峰底到峰顶, 再找到峰底到封顶. 循环往复!
        """
        self.len = len(prices)
        if self.len <= 1:
            return 0

        self.prices = prices

        profit = 0

        start_idx = 0
        while True:
            buy_idx = self.find_min_price(start_idx)
            if buy_idx == -1:
                break

            # print('buy ', buy_idx)
            sell_idx = self.find_max_price(buy_idx)
            if sell_idx == -1:
                break

            # print('sell ', sell_idx)
            profit += self.prices[sell_idx] - self.prices[buy_idx]
            start_idx = sell_idx

        return profit

    def find_min_price(self, from_idx):
        # NOTE: valley
        pre_idx = from_idx

        for next_idx in range(from_idx+1, self.len):
            if self.prices[next_idx] > self.prices[pre_idx]:
                return pre_idx
            else:
                pre_idx += 1

        return -1

    def find_max_price(self, from_idx):
        # NOTE: peak
        pre_idx = from_idx

        next_idx = -1
        for next_idx in range(from_idx+1, self.len):
            if self.prices[next_idx] < self.prices[pre_idx]:
                # print("m|return ", pre_idx)
                return pre_idx
            else:
                pre_idx += 1
                # print("after add ", pre_idx)

        assert next_idx != -1
        # print("f|return ", next_idx)
        return next_idx


    def method1(self, prices):
        # 动态规划法
        if not prices:
            return 0

        self.prices = prices
        self.len = len(prices)
        self.memory = {}

        return self.max_profit_rec(0, self.len)

    def max_profit_rec(self, b, e):
        # [b, e), 开闭区间
        key = b, e
        max_pro = self.memory.get(key)

        # print(key)

        if max_pro is not None:
            return max_pro

        if e - b == 1:
            max_pro = 0
        else:
            max_pro = max(0, self.prices[e-1] - self.prices[b])
            for mid in range(b+1, e):
                pro = self.max_profit_rec(b, mid) + self.max_profit_rec(mid, e)
                max_pro = max(pro, max_pro)

        self.memory[key] = max_pro
        # print(key, max_pro)
        return max_pro


def test():
    s = Solution()

    lst = [2, 1, 2, 0, 1]
    assert s.maxProfit(lst) == 2

    lst = [0, 1]
    assert s.maxProfit(lst) == 1

    lst = [4, 3, 2, 1]
    assert s.maxProfit(lst) == 0

def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
